from typing import List, Dict, Any, Optional
import logging
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from .base import (
    BaseGraphAugmenter,
    BaseGraphReasoner,
    Node,
    Edge,
    SubGraph
)


class GraphAugmenter(BaseGraphAugmenter):
    """图增强实现"""
    
    def __init__(self):
        self.llm = None
        self.logger = logging.getLogger(__name__)
        
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化增强器"""
        self.llm = ChatOpenAI(
            model_name=config.get("model_name", "gpt-3.5-turbo"),
            temperature=config.get("temperature", 0),
            openai_api_key=config["openai_api_key"]
        )
        
    async def augment(self, query: str, subgraph: SubGraph) -> SubGraph:
        """增强子图"""
        # 构建提示模板
        template = """基于以下查询和子图信息，请推断可能存在的隐含关系和属性。

查询: {query}

子图信息:
节点:
{nodes}

边:
{edges}

请分析并推断:
1. 节点之间可能存在的隐含关系
2. 节点可能具有的额外属性
3. 子图中缺失的关键信息

请以JSON格式输出，包含以下字段：
- new_nodes: 新增节点列表
- new_edges: 新增边列表
- node_updates: 节点属性更新列表

输出格式示例：
{
    "new_nodes": [
        {"id": "n1", "type": "概念", "properties": {"name": "示例概念"}}
    ],
    "new_edges": [
        {"source": "已有节点id", "target": "n1", "type": "相关", "properties": {}}
    ],
    "node_updates": [
        {"id": "已有节点id", "properties": {"新属性": "值"}}
    ]
}

请仅输出JSON，不要包含其他说明文本。"""
        
        # 准备输入
        nodes_text = "\n".join(
            f"- ID: {node.id}, 类型: {node.type}, 属性: {node.properties}"
            for node in subgraph.nodes
        )
        
        edges_text = "\n".join(
            f"- {edge.source_id} --[{edge.type}]--> {edge.target_id}"
            for edge in subgraph.edges
        )
        
        # 创建提示
        prompt = ChatPromptTemplate.from_template(template)
        
        try:
            # 调用LLM
            response = await self.llm.ainvoke(
                prompt.format_messages(
                    query=query,
                    nodes=nodes_text,
                    edges=edges_text
                )
            )
            
            # 解析响应
            augmentation = response.content
            if not isinstance(augmentation, dict):
                import json
                augmentation = json.loads(augmentation)
            
            # 应用增强
            augmented_subgraph = self._apply_augmentation(subgraph, augmentation)
            return augmented_subgraph
            
        except Exception as e:
            self.logger.error(f"图增强过程出错: {str(e)}")
            return subgraph
            
    def _apply_augmentation(self, subgraph: SubGraph, augmentation: Dict[str, Any]) -> SubGraph:
        """应用增强结果到子图"""
        # 创建新的子图对象
        new_subgraph = SubGraph(
            nodes=list(subgraph.nodes),
            edges=list(subgraph.edges)
        )
        
        # 添加新节点
        for node_data in augmentation.get("new_nodes", []):
            new_node = Node(
                id=node_data["id"],
                type=node_data["type"],
                properties=node_data["properties"]
            )
            new_subgraph.nodes.append(new_node)
            
        # 添加新边
        for edge_data in augmentation.get("new_edges", []):
            new_edge = Edge(
                source_id=edge_data["source"],
                target_id=edge_data["target"],
                type=edge_data["type"],
                properties=edge_data["properties"]
            )
            new_subgraph.edges.append(new_edge)
            
        # 更新节点属性
        for update in augmentation.get("node_updates", []):
            node_id = update["id"]
            for node in new_subgraph.nodes:
                if node.id == node_id:
                    node.properties.update(update["properties"])
                    break
                    
        return new_subgraph


class GraphReasoner(BaseGraphReasoner):
    """图推理实现"""
    
    def __init__(self):
        self.llm = None
        self.logger = logging.getLogger(__name__)
        
    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化推理器"""
        self.llm = ChatOpenAI(
            model_name=config.get("model_name", "gpt-3.5-turbo"),
            temperature=config.get("temperature", 0),
            openai_api_key=config["openai_api_key"]
        )
        
    async def reason(self, query: str, subgraph: SubGraph) -> Dict[str, Any]:
        """执行推理"""
        # 构建提示模板
        template = """基于以下查询和知识图谱信息，请进行推理并回答问题。

查询: {query}

知识图谱信息:
节点:
{nodes}

关系:
{edges}

请提供以下内容：
1. 详细的推理过程
2. 最终答案
3. 使用的关键证据
4. 可信度评估

请以JSON格式输出，包含以下字段：
- reasoning_steps: 推理步骤列表
- answer: 最终答案
- evidence: 使用的关键证据列表
- confidence: 可信度分数(0-1)
- uncertainty: 不确定性说明

输出格式示例：
{
    "reasoning_steps": [
        "首先分析...",
        "然后发现...",
        "最后得出..."
    ],
    "answer": "详细的答案",
    "evidence": [
        "证据1: ...",
        "证据2: ..."
    ],
    "confidence": 0.85,
    "uncertainty": "可能存在的不确定因素..."
}

请仅输出JSON，不要包含其他说明文本。"""
        
        # 准备输入
        nodes_text = "\n".join(
            f"- ID: {node.id}, 类型: {node.type}, 属性: {node.properties}"
            for node in subgraph.nodes
        )
        
        edges_text = "\n".join(
            f"- {edge.source_id} --[{edge.type}]--> {edge.target_id}"
            for edge in subgraph.edges
        )
        
        # 创建提示
        prompt = ChatPromptTemplate.from_template(template)
        
        try:
            # 调用LLM
            response = await self.llm.ainvoke(
                prompt.format_messages(
                    query=query,
                    nodes=nodes_text,
                    edges=edges_text
                )
            )
            
            # 解析响应
            reasoning_result = response.content
            if not isinstance(reasoning_result, dict):
                import json
                reasoning_result = json.loads(reasoning_result)
                
            return reasoning_result
            
        except Exception as e:
            self.logger.error(f"图推理过程出错: {str(e)}")
            return {
                "reasoning_steps": ["推理过程出错"],
                "answer": "无法完成推理",
                "evidence": [],
                "confidence": 0.0,
                "uncertainty": str(e)
            } 